#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/8/12 11:02
# @Author  : Tingkyen
# @File    : demo_grid_polygon.py
# @Software: PyCharm
import time

import numpy as np
import shapely
from shapely.geometry import Polygon, box
from shapely.strtree import STRtree

def discretize_coordinates(coordinates,min_x,max_x,min_y,max_y,grid_size=128):
    discretize_coords = []
    grid_width = (max_x - min_x) / grid_size
    grid_height = (max_y - min_y) / grid_size

    for (x,y) in coordinates:
        i = int(np.floor((x - min_x) / grid_width))
        j = int(np.floor((y-min_y) / grid_height))
        if i >= grid_size:
            i = grid_size - 1
        if j >= grid_size:
            j = grid_size - 1
        discretize_coords.append((i,j))
    return discretize_coords


def discretize_polygon(polygon, grid_size=128):
    """
    将polygon离散化
    将多边形离散化到 grid_size x grid_size 网格中
    支持两种模式：
        - 'binary':  是否相交（0/1）
        - 'area_ratio': 按重叠面积比例赋值 (0~1)
    :param polygon: 例Polygon([(0.5, 0.7), (1.2, 2.3), (3.4, 4.5), (0.5, 0.7)])
    :param grid_size:
    :return: 2维数组 (grid_size, grid_size)
    """
    # 定义空间范围和网格大小
    min_x, min_y, max_x, max_y = polygon.bounds

    # 计算网格单元尺寸
    grid_width = (max_x - min_x) / grid_size
    grid_height = (max_y - min_y) / grid_size

    # 快速退出：如果多边形为空或退化
    if grid_width == 0 or grid_height == 0:
        return np.zeros((grid_size, grid_size))

    # 初始化网格
    grid = np.zeros((grid_size, grid_size))

    # 遍历所有网格单元，检查是否与多边形相交
    for i in range(grid_size):
        for j in range(grid_size):
            # 计算当前网格单元的边界框
            cell_bounds = box(min_x + i * grid_width,
                              min_y + j * grid_height,
                              min_x + (i + 1) * grid_width,
                              min_y + (j + 1) * grid_height)

            # 判断是否与多边形相交
            if polygon.intersects(cell_bounds):
                # 可以根据需要调整这里的逻辑
                # 这里简单的标记为1表示相交
                grid[j, i] = 1  # 注意numpy数组索引顺序[y, x]

    return grid


def discretize_polygon_pro(polygon, grid_size=128, method='area_ratio'):
    """
    将多边形离散化到 grid_size x grid_size 网格中
    支持两种模式：
        - 'binary':  是否相交（0/1）
        - 'area_ratio': 按重叠面积比例赋值 (0~1)

    :param polygon: shapely Polygon 对象
    :param grid_size: 网格边长，如 128
    :param method: 'binary' 或 'area_ratio'
    :return: 2D numpy array (grid_size, grid_size)
    """
    min_x, min_y, max_x, max_y = polygon.bounds

    # 计算网格单元尺寸
    width = (max_x - min_x) / grid_size
    height = (max_y - min_y) / grid_size

    # 快速退出：如果多边形为空或退化
    if width == 0 or height == 0:
        return np.zeros((grid_size, grid_size))

    # Step 1: 构建候选网格索引（仅检查包围盒相交的网格）
    # 使用网格化方式生成所有网格的 bounds（向量化）
    cols = np.arange(grid_size)
    rows = np.arange(grid_size)

    # 向量化生成所有左上角坐标
    x_coords = min_x + np.outer(np.ones(grid_size), cols) * width  # (128, 128)
    y_coords = min_y + np.outer(rows, np.ones(grid_size)) * height

    # 所有网格的 bounds: (left, bottom, right, top)
    left = x_coords
    bottom = y_coords
    right = left + width
    top = bottom + height

    # Step 2: 使用 STRtree 加速空间查询
    # 创建所有网格的 box bounds 列表（只创建可能相交的）
    candidate_boxes = []
    candidate_indices = []

    # 先用包围盒粗筛（更快）
    poly_bounds = polygon.bounds  # (minx, miny, maxx, maxy)

    for i in range(grid_size):
        for j in range(grid_size):
            cell_right = left[i, j] + width
            cell_top = bottom[i, j] + height
            if (left[i, j] < poly_bounds[2] and cell_right > poly_bounds[0] and
                bottom[i, j] < poly_bounds[3] and cell_top > poly_bounds[1]):

                cell_box = box(left[i, j], bottom[i, j], cell_right, cell_top)
                candidate_boxes.append(cell_box)
                candidate_indices.append((i, j))

    if not candidate_boxes:
        return np.zeros((grid_size, grid_size))

    # 构建 STRtree 提高 intersects 查询效率
    tree = STRtree(candidate_boxes)

    # 查询与 polygon 相交的所有候选网格
    intersecting = tree.query(polygon, predicate='intersects')

    # 初始化结果网格
    grid = np.zeros((grid_size, grid_size))

    # 只处理真正相交的网格
    for idx in intersecting:
        i, j = candidate_indices[idx]
        cell = candidate_boxes[idx]

        if method == 'binary':
            grid[j, i] = 1  # 注意：[j, i] 对应 [col, row]，但存储为 [y, x]

        elif method == 'area_ratio':
            intersection = polygon.intersection(cell)
            cell_area = width * height
            grid[j, i] = intersection.area / cell_area if cell_area > 0 else 0

    return grid


def discretize_polygon_vectorized(polygon, grid_size=128, method='area_ratio'):
    """
    使用 shapely 2.0 向量化操作加速
    """
    min_x, min_y, max_x, max_y = shapely.bounds(polygon)
    width = (max_x - min_x) / grid_size
    height = (max_y - min_y) / grid_size

    if width == 0 or height == 0:
        return np.zeros((grid_size, grid_size))

    # 向量化生成所有网格的 corners
    cols = np.arange(grid_size)
    rows = np.arange(grid_size)
    x_left = min_x + cols * width
    y_bottom = min_y + rows * height

    # 构建所有 grid cells（一次性向量化创建）
    left, bottom = np.meshgrid(x_left, y_bottom, indexing='xy')  # (grid_size, grid_size)
    right = left + width
    top = bottom + height

    # 向量化创建所有 boxes
    cells = shapely.box(left.ravel(), bottom.ravel(), right.ravel(), top.ravel())

    # 向量化判断是否相交
    intersects = shapely.intersects(polygon, cells)  # (grid_size*grid_size)

    # 初始化结果
    grid = np.zeros(grid_size * grid_size)

    if method == 'binary':
        grid[intersects] = 1
    elif method == 'area_ratio':
        # 向量化求交集面积（注意：polygon 是 scalar，cells 是 array）
        intersections = shapely.intersection(polygon, cells[intersects])
        cell_area = width * height
        grid[intersects] = shapely.area(intersections) / cell_area

    return grid.reshape((grid_size, grid_size))


def print_hi(name):
    print(f"Hi, {name}")


def test_discretize():
    coordinates = [(0.5, 0.7), (1.2, 2.3), (3.4, 4.5)]
    discretized = discretize_coordinates(coordinates, min_x=0, max_x=5, min_y=0, max_y=5, grid_size=128)
    print(discretized)
    # 示例多边形定义
    polygon = Polygon([(0.5, 0.7), (1.2, 2.3), (3.4, 4.5), (0.5, 0.7)])
    t1 = time.time()
    discretize_polygon(polygon,512)
    # 计算运行时间
    t2 = time.time()
    print(f"discretize_polygon(polygon,512)运行时间：{t2-t1}")
    discretize_polygon_pro(polygon,512,"binary")
    t3 = time.time()
    print(f"discretize_polygon_pro(polygon,512,\"binary\")运行时间：{t3-t2}")
    discretize_polygon_vectorized(polygon,512,"binary")
    t4 = time.time()
    print(f"discretize_polygon_vectorized运行时间：{t4-t3}")


if __name__ == "__main__":
    # test_discretize()
    polygon = Polygon([(0.5, 0.7), (1.2, 2.3), (3.4, 4.5), (0.5, 0.7)])
    discretize_polygon_vectorized(polygon,512,"binary")
    print_hi("Tingkyen")
